Downstram Analysis of the results of the Deconvolution

import sys
import pickle
import scanpy as sc
import anndata
import numpy as np
import os
import pandas as pd
import seaborn as sns

import cell2location

import matplotlib as mpl
from matplotlib import pyplot as plt

# scanpy prints a lot of warnings
import warnings
warnings.filterwarnings('ignore')

results_folder = 'Cell2Location/results/'
r = {'run_name': 'LocationModelLinearDependentWMultiExperiment_14experiments_36clusters_20654locations_4188genes'}

# defining useful function
def select_slide(adata, s, s_col='sample'):
    r""" Select data for one slide from the spatial anndata object.

    :param adata: Anndata object with multiple spatial samples
    :param s: name of selected sample
    :param s_col: column in adata.obs listing sample name for each location
    """

    slide = adata[adata.obs[s_col].isin([s]), :]
    s_keys = list(slide.uns['spatial'].keys())
    s_spatial = np.array(s_keys)[[s in k for k in s_keys]][0]

    slide.uns['spatial'] = {s_spatial: slide.uns['spatial'][s_spatial]}

    return slide

sp_data_file = results_folder +r['run_name']+'/sp.h5ad'

adata_vis = anndata.read(sp_data_file)

adata_vis.obs

	in_tissue	array_row	array_col	sample	n_genes_by_counts	log1p_n_genes_by_counts	total_counts	log1p_total_counts	pct_counts_in_top_50_genes	pct_counts_in_top_100_genes	...	q95_nUMI_factorsStem-like/TA	q95_nUMI_factorsStromal 1	q95_nUMI_factorsStromal 2	q95_nUMI_factorsStromal 3	q95_nUMI_factorsT follicular helper cells	q95_nUMI_factorsT helper 17 cells	q95_nUMI_factorsTip-like ECs	q95_nUMI_factorsUnknown	q95_nUMI_factorscDC	q95_nUMI_factorsgamma delta T cells
spot_id
Count_SN048_A121573_Rep1_AAACAAGTATCTCCCA-1	1	50	102	Count_SN048_A121573_Rep1	5341	8.583355	17364.0	9.762212	25.662290	33.920756	...	1034.578638	66.694864	25.468667	3.680339	8.581451	6.973702	73.191489	23.585139	18.172732	8.738241
Count_SN048_A121573_Rep1_AAACACCAATAACTGC-1	1	59	19	Count_SN048_A121573_Rep1	5674	8.643826	22320.0	10.013283	28.369176	38.342294	...	73.807870	8.363901	40.141215	4.890074	17.453783	13.598687	31.092448	24.300600	24.653172	9.175505
Count_SN048_A121573_Rep1_AAACATTTCCCGGATT-1	1	61	97	Count_SN048_A121573_Rep1	3170	8.061802	7669.0	8.945072	24.149172	33.224671	...	172.479327	15.275673	12.273451	3.026998	7.953982	5.950548	21.504271	20.162226	9.718825	5.554598
Count_SN048_A121573_Rep1_AAACCCGAACGAAATC-1	1	45	115	Count_SN048_A121573_Rep1	6756	8.818334	27100.0	10.207326	19.712177	27.811808	...	211.151180	10.889569	54.313984	4.840832	13.695137	16.807419	20.052122	26.854432	20.586775	11.660860
Count_SN048_A121573_Rep1_AAACCGGGTAGGTACC-1	1	42	28	Count_SN048_A121573_Rep1	4295	8.365440	12265.0	9.414586	23.742356	32.768039	...	53.288616	3.379716	19.922063	13.277920	11.294168	7.806391	7.958049	20.324797	6.271026	4.446291
...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...
SN123_A938797_Rep1_TTGTTAGCAAATTCGA-1	1	22	42	SN123_A938797_Rep1	2339	7.757906	3494.0	8.159089	12.850601	19.748140	...	12.410881	48.224864	7.234528	5.231929	40.203137	16.265775	3.136449	13.240821	5.791877	8.011834
SN123_A938797_Rep1_TTGTTCAGTGTGCTAC-1	1	24	64	SN123_A938797_Rep1	3280	8.095904	6766.0	8.819813	17.188886	25.022170	...	30.029537	10.171256	61.009798	9.339274	5.686483	5.889894	27.038523	16.841162	30.875569	4.636048
SN123_A938797_Rep1_TTGTTGTGTGTCAAGA-1	1	31	77	SN123_A938797_Rep1	4068	8.311153	9349.0	9.143131	17.574072	25.414483	...	44.019627	17.114798	47.078424	48.573906	13.218768	9.802032	20.935667	23.153979	20.348671	6.516171
SN123_A938797_Rep1_TTGTTTCACATCCAGG-1	1	58	42	SN123_A938797_Rep1	3365	8.121480	6868.0	8.834774	17.850903	24.548631	...	37.958709	33.372944	311.826648	41.852420	14.511507	12.165736	146.332658	40.171212	30.095921	9.832882
SN123_A938797_Rep1_TTGTTTCCATACAACT-1	1	45	27	SN123_A938797_Rep1	2805	7.939515	4937.0	8.504716	16.021876	23.617582	...	27.401279	43.675632	74.343894	26.139359	27.500353	20.165950	57.000325	42.849142	55.225842	11.603575

20654 rows × 301 columns

1.- Clustering based on cell abundances

Low Resolution = 0.25 (Few clusters)

sample_type = 'q05_nUMI_factors'
col_ind = [sample_type in i for i in adata_vis.obs.columns.tolist()]
adata_vis.obsm[sample_type] = adata_vis.obs.loc[:,col_ind].values

# compute KNN using the cell2location output
sc.pp.neighbors(adata_vis, use_rep=sample_type,
                n_neighbors = 30)

# Cluster spots into regions using scanpy
sc.tl.leiden(adata_vis, resolution=0.25)

# add region as categorical variable
adata_vis.obs["region_cluster_025"] = adata_vis.obs["leiden"]
adata_vis.obs["region_cluster_025"] =  adata_vis.obs["region_cluster_025"].astype("category")

sc.tl.umap(adata_vis, min_dist = 0.75, spread = 1)

with mpl.rc_context({'figure.figsize': (16, 16)}):
    sc.pl.umap(adata_vis, color=['sample'], size=45,
               color_map = 'RdPu',
               legend_fontsize=12, alpha=0.75)

with mpl.rc_context({'figure.figsize': (16, 16)}):
    sc.pl.umap(adata_vis, color=['region_cluster_025'], size=45,
               color_map = 'RdPu',
               legend_fontsize=12, alpha=0.75)

with mpl.rc_context({'figure.figsize': (12, 12)}):
  sc.pl.spatial(adata_vis[adata_vis.obs["sample"]=="Count_SN123_A595688_Rep1"],
                  color=["region_cluster_025"], size= 0.15, img_key=None
                )

Higher Resolution = 0.75 (More clusters)

sample_type = 'q05_nUMI_factors'
col_ind = [sample_type in i for i in adata_vis.obs.columns.tolist()]
adata_vis.obsm[sample_type] = adata_vis.obs.loc[:,col_ind].values

# compute KNN using the cell2location output
sc.pp.neighbors(adata_vis, use_rep=sample_type,
                n_neighbors = 30)

# Cluster spots into regions using scanpy
sc.tl.leiden(adata_vis, resolution=0.75)

# add region as categorical variable
adata_vis.obs["region_cluster_075"] = adata_vis.obs["leiden"]
adata_vis.obs["region_cluster_075"] =  adata_vis.obs["region_cluster_075"].astype("category")

sc.tl.umap(adata_vis, min_dist = 0.75, spread = 1)

with mpl.rc_context({'figure.figsize': (16, 16)}):
    sc.pl.umap(adata_vis, color=['sample'], size=45,
               color_map = 'RdPu',
               legend_fontsize=12, alpha=0.75)

with mpl.rc_context({'figure.figsize': (16, 16)}):
    sc.pl.umap(adata_vis, color=['region_cluster_075'], size=45,
               color_map = 'RdPu',
               legend_fontsize=12, alpha=0.75)

with mpl.rc_context({'figure.figsize': (12, 12)}):
  sc.pl.spatial(adata_vis[adata_vis.obs["sample"]=="Count_SN123_A595688_Rep1"],
                  color=["region_cluster_075"], size= 0.15, img_key=None
                )

We save the results to compute the Neigborhod enrichment analysis in the following script.

sp_data_New_file = results_folder +r['run_name']+'/sp_clusters_proportions.h5ad'
adata_vis.write(filename=sp_data_New_file)

2.- Proportions of all the cell types per sample-

samples = {'sample_name': ['Count_SN048_A121573_Rep1', 'Count_SN048_A121573_Rep2', 'Count_SN048_A416371_Rep1', 
    'Count_SN048_A416371_Rep2','Count_SN123_A551763_Rep1', 'Count_SN123_A595688_Rep1', 'Count_SN123_A798015_Rep1', 
    'Count_SN124_A551763_Rep2','Count_SN124_A595688_Rep2','Count_SN124_A798015_Rep2',
    'Count_SN124_A938797_Rep2','Count_SN84_A120838_Rep1','Count_SN84_A120838_Rep2', 'SN123_A938797_Rep1']}
sample_data = pd.DataFrame(data=samples)
sample_type = 'q05_nUMI_factors'
col_ind = [sample_type in i for i in adata_vis.obs.columns.tolist()]
df=pd.DataFrame()

for i in sample_data['sample_name']:
    adata_sample = adata_vis.obs.query('sample == @i')
    adata_sample_celltypes = adata_sample.loc[:,col_ind]
    per_sample_percelltype_counts = adata_sample_celltypes.sum()
    per_sample_totalcounts = per_sample_percelltype_counts.sum()
    per_sample_percelltype_proportions = per_sample_percelltype_counts/per_sample_totalcounts
    df[i] = per_sample_percelltype_proportions
   

df

	Count_SN048_A121573_Rep1	Count_SN048_A121573_Rep2	Count_SN048_A416371_Rep1	Count_SN048_A416371_Rep2	Count_SN123_A551763_Rep1	Count_SN123_A595688_Rep1	Count_SN123_A798015_Rep1	Count_SN124_A551763_Rep2	Count_SN124_A595688_Rep2	Count_SN124_A798015_Rep2	Count_SN124_A938797_Rep2	Count_SN84_A120838_Rep1	Count_SN84_A120838_Rep2	SN123_A938797_Rep1
q05_nUMI_factorsCD19+CD20+ B	0.002734	0.002391	0.125460	0.129224	0.042431	0.002793	0.006635	0.042895	0.000953	0.005268	0.003613	0.001611	0.001878	0.005901
q05_nUMI_factorsCD4+ T cells	0.000563	0.000640	0.000928	0.001044	0.001248	0.001193	0.001671	0.001419	0.000896	0.001350	0.002205	0.001246	0.001190	0.002234
q05_nUMI_factorsCD8+ T cells	0.000614	0.000572	0.025399	0.026373	0.007983	0.000561	0.001300	0.008295	0.000261	0.001058	0.000935	0.000376	0.000429	0.001081
q05_nUMI_factorsCMS1	0.003255	0.003933	0.442475	0.404469	0.278569	0.001020	0.022809	0.333333	0.000770	0.017643	0.003982	0.000933	0.001490	0.002389
q05_nUMI_factorsCMS2	0.383587	0.470242	0.216457	0.192636	0.254264	0.418494	0.011939	0.193486	0.576922	0.014241	0.309458	0.614015	0.675785	0.279049
q05_nUMI_factorsCMS3	0.081321	0.068295	0.012452	0.019755	0.041942	0.021303	0.052904	0.083533	0.010600	0.075666	0.010581	0.000211	0.000697	0.004956
q05_nUMI_factorsCMS4	0.000531	0.000638	0.001375	0.001497	0.000717	0.001010	0.000721	0.000764	0.000778	0.000642	0.001304	0.000829	0.000800	0.001150
q05_nUMI_factorsEnteric glial cells	0.000525	0.000339	0.000127	0.000131	0.000335	0.000822	0.003432	0.000682	0.000538	0.002398	0.002212	0.000574	0.000764	0.002391
q05_nUMI_factorsGoblet cells	0.068542	0.052916	0.005560	0.009539	0.011276	0.058924	0.059073	0.007807	0.059544	0.081506	0.012783	0.000106	0.000343	0.014333
q05_nUMI_factorsIgA+ Plasma	0.075952	0.060537	0.000861	0.001427	0.000645	0.037327	0.032408	0.000551	0.017404	0.072281	0.019883	0.000363	0.000409	0.012754
q05_nUMI_factorsIgG+ Plasma	0.018209	0.021158	0.001301	0.001875	0.010451	0.028145	0.202003	0.004521	0.024616	0.181026	0.106688	0.010481	0.021169	0.105342
q05_nUMI_factorsIntermediate	0.001693	0.001228	0.000335	0.000389	0.000235	0.001040	0.001384	0.000199	0.000390	0.002061	0.000733	0.000230	0.000226	0.000578
q05_nUMI_factorsLymphatic ECs	0.000654	0.000648	0.000346	0.000468	0.001766	0.001436	0.003074	0.001128	0.001377	0.002513	0.002822	0.001682	0.001276	0.002690
q05_nUMI_factorsMast cells	0.000306	0.000258	0.000097	0.000128	0.000647	0.000682	0.002129	0.000536	0.000402	0.001705	0.001450	0.000446	0.000429	0.001435
q05_nUMI_factorsMature Enterocytes type 1	0.028608	0.022936	0.002158	0.002294	0.001502	0.012936	0.019966	0.000859	0.003710	0.027783	0.006833	0.000490	0.000642	0.006857
q05_nUMI_factorsMature Enterocytes type 2	0.154321	0.131087	0.013281	0.014353	0.017378	0.072557	0.109520	0.018424	0.023303	0.142275	0.034494	0.002936	0.004622	0.030453
q05_nUMI_factorsMyofibroblasts	0.010301	0.016013	0.030209	0.043955	0.190093	0.112392	0.115479	0.159551	0.096002	0.062894	0.263073	0.208643	0.137958	0.322069
q05_nUMI_factorsNK cells	0.000222	0.000213	0.003373	0.003707	0.000999	0.000295	0.000433	0.000756	0.000168	0.000400	0.000517	0.000227	0.000220	0.000467
q05_nUMI_factorsPericytes	0.006044	0.008659	0.010607	0.014643	0.029540	0.019884	0.017048	0.019860	0.017304	0.009938	0.025786	0.033568	0.029030	0.027667
q05_nUMI_factorsPro-inflammatory	0.002140	0.003007	0.002360	0.002460	0.006845	0.002754	0.001497	0.026106	0.001391	0.001562	0.001741	0.002778	0.001776	0.001485
q05_nUMI_factorsProliferating	0.004443	0.005926	0.010818	0.011577	0.014032	0.007272	0.001089	0.020154	0.004720	0.000904	0.007227	0.005649	0.004239	0.004946
q05_nUMI_factorsProliferative ECs	0.004034	0.005272	0.012177	0.015023	0.026723	0.006780	0.000729	0.023252	0.004535	0.000685	0.006323	0.014094	0.008982	0.004469
q05_nUMI_factorsRegulatory T cells	0.000556	0.000611	0.001285	0.001411	0.001377	0.001006	0.001354	0.001367	0.000716	0.001144	0.001886	0.000986	0.000928	0.001926
q05_nUMI_factorsSPP1+	0.004049	0.005925	0.050844	0.065007	0.021440	0.010735	0.022768	0.019664	0.008510	0.017403	0.034586	0.026193	0.020249	0.025179
q05_nUMI_factorsSmooth muscle cells	0.001931	0.002290	0.000657	0.001022	0.005464	0.013652	0.078866	0.004593	0.015579	0.058344	0.018802	0.008996	0.013507	0.022005
q05_nUMI_factorsStalk-like ECs	0.002237	0.002064	0.000580	0.000782	0.003582	0.003118	0.016124	0.004333	0.002080	0.011196	0.006523	0.011620	0.007482	0.007327
q05_nUMI_factorsStem-like/TA	0.122502	0.092340	0.011516	0.012895	0.007109	0.093789	0.072831	0.004597	0.076007	0.100790	0.031586	0.002405	0.002014	0.027733
q05_nUMI_factorsStromal 1	0.004392	0.002970	0.000715	0.000783	0.000493	0.004940	0.004006	0.000766	0.002617	0.006259	0.008300	0.000744	0.001239	0.008446
q05_nUMI_factorsStromal 2	0.008481	0.009400	0.000649	0.000572	0.004164	0.038235	0.005755	0.003186	0.029868	0.009538	0.023273	0.004086	0.006470	0.023382
q05_nUMI_factorsStromal 3	0.000471	0.000399	0.000211	0.000305	0.001833	0.010208	0.111257	0.001153	0.008711	0.071057	0.027062	0.030781	0.038556	0.027719
q05_nUMI_factorsT follicular helper cells	0.000474	0.000512	0.001592	0.001732	0.001076	0.000904	0.001121	0.001017	0.000656	0.000922	0.001634	0.000865	0.000821	0.001653
q05_nUMI_factorsT helper 17 cells	0.000452	0.000464	0.001271	0.001412	0.000866	0.000742	0.000952	0.000735	0.000493	0.000867	0.001269	0.000613	0.000577	0.001236
q05_nUMI_factorsTip-like ECs	0.003711	0.004144	0.001568	0.002357	0.007222	0.009235	0.010790	0.006678	0.006330	0.008123	0.009759	0.008602	0.009910	0.011191
q05_nUMI_factorsUnknown	0.000408	0.000463	0.002017	0.002467	0.000869	0.000994	0.002060	0.000503	0.000515	0.001440	0.002332	0.000744	0.000751	0.002019
q05_nUMI_factorscDC	0.001391	0.001182	0.007935	0.011158	0.004435	0.002394	0.004354	0.002963	0.001090	0.006578	0.007653	0.001558	0.002845	0.004859
q05_nUMI_factorsgamma delta T cells	0.000349	0.000331	0.001004	0.001129	0.000447	0.000428	0.000517	0.000333	0.000244	0.000536	0.000692	0.000320	0.000300	0.000627

g = sns.clustermap(df,figsize=(18, 14),row_cluster=False,cmap="vlag", metric="correlation", z_score=0)

3-. Identify groups of co-located cell types using matrix factorisation

From Cell2Location tutorial: Tip If you want to find a few most disctinct cellular compartments, use a small number of factors. If you want to find very strong co-location signal and assume that most cell types don’t co-locate, use a lot of factors (> 30 - used here). In practice, it is better to train NMF for a range of factors R=5,..,30 and select R as a balance between capturing fine tissue zones and splitting known compartments

R = 5

# number of cell type combinations - educated guess assuming that most cell types don't co-locate
n_fact = int(5)

# extract cell abundance from cell2location
X_data = adata_vis.uns['mod']['post_sample_q05']['spot_factors']

import cell2location.models as c2l
# create model class
mod_sk = c2l.CoLocatedGroupsSklearnNMF(n_fact, X_data,
        n_iter = 10000,
        verbose = True,
        var_names=adata_vis.uns['mod']['fact_names'],
        obs_names=adata_vis.obs_names,
        fact_names=['fact_' + str(i) for i in range(n_fact)],
        sample_id=adata_vis.obs['sample'],
        init='random', random_state=0,
        nmf_kwd_args={'tol':0.0001})

# train 5 times to evaluate stability
mod_sk.fit(n=5, n_type='restart')

init_1 - iterations until convergence: 6517
init_2 - iterations until convergence: 6568
init_3 - iterations until convergence: 6693
init_4 - iterations until convergence: 7243
init_5 - iterations until convergence: 5590

## Do some diagnostics
# evaluate stability by comparing trainin restarts

with mpl.rc_context({'figure.figsize': (10, 8)}):
    mod_sk.evaluate_stability('cell_type_factors', align=True)

# evaluate accuracy of the model
mod_sk.compute_expected()
mod_sk.plot_posterior_mu_vs_data()

# extract parameters into DataFrames
mod_sk.sample2df(node_name='nUMI_factors', ct_node_name = 'cell_type_factors')

# export results to scanpy object
adata_vis = mod_sk.annotate_adata(adata_vis) # as columns to .obs
adata_vis = mod_sk.export2adata(adata_vis, slot_name='mod_sklearn') # as a slot in .uns

# print the fraction of cells of each type located to each combination
mod_sk.print_gene_loadings(loadings_attr='cell_type_fractions',
                         gene_fact_name='cell_type_fractions', top_n = 9)

	top-1	top-2	top-3	top-4	top-5	top-6	top-7	top-8	top-9
mean_cell_type_factorsfact_0	Mature Enterocytes type 2: 0.98	Mature Enterocytes type 1: 0.97	Goblet cells: 0.97	Stem-like/TA: 0.87	CMS3: 0.78	IgA+ Plasma: 0.65	Intermediate: 0.61	Stromal 1: 0.17	Stromal 2: 0.091
mean_cell_type_factorsfact_1	CMS2: 0.96	Stromal 2: 0.44	Pro-inflammatory: 0.38	Proliferative ECs: 0.37	Proliferating: 0.37	CMS4: 0.24	Pericytes: 0.2	T follicular helper cells: 0.14	CD4+ T cells: 0.13
mean_cell_type_factorsfact_2	Stromal 3: 1.0	IgG+ Plasma: 1.0	Smooth muscle cells: 0.99	Myofibroblasts: 0.98	Stalk-like ECs: 0.95	Mast cells: 0.94	Enteric glial cells: 0.91	Lymphatic ECs: 0.9	Tip-like ECs: 0.79
mean_cell_type_factorsfact_3	CMS1: 0.98	CD19+CD20+ B: 0.53	CMS3: 0.17	Pro-inflammatory: 0.14	CMS4: 0.12	Unknown: 0.092	cDC: 0.07	CMS2: 0.04	Enteric glial cells: 0.031
mean_cell_type_factorsfact_4	CD8+ T cells: 0.98	NK cells: 0.75	CD19+CD20+ B: 0.47	Proliferative ECs: 0.39	SPP1+: 0.39	Proliferating: 0.35	gamma delta T cells: 0.35	cDC: 0.33	T helper 17 cells: 0.27

# make nice names
from re import sub
mod_sk.cell_type_fractions.columns = [sub('mean_cell_type_factors', '', i)
                                      for i in mod_sk.cell_type_fractions.columns]

# plot co-occuring cell type combinations
mod_sk.plot_gene_loadings(mod_sk.var_names_read, mod_sk.var_names_read,
                        fact_filt=mod_sk.fact_filt,
                        loadings_attr='cell_type_fractions',
                        gene_fact_name='cell_type_fractions',
                        cmap='RdPu', figsize=[10, 10])

# plot cell density in each combination
with mpl.rc_context({'figure.figsize': (8, 8), 'axes.facecolor': 'black'}):

    # select one section correctly subsetting histology image data
    slide = select_slide(adata_vis, 'Count_SN048_A121573_Rep2')

    sc.pl.spatial(slide,
                  cmap='viridis',
                  color=mod_sk.location_factors_df.columns,
                  size=1, img_key='hires',
                  alpha_img=0.5, alpha = 0.4, ncols=2,
                  vmin=0, vmax='p99.2'
                 )

Trying to set attribute `.uns` of view, copying.

# plot cell density in each combination
with mpl.rc_context({'figure.figsize': (8, 8), 'axes.facecolor': 'black'}):

    # select one section correctly subsetting histology image data
    slide = select_slide(adata_vis, 'Count_SN048_A416371_Rep1')

    sc.pl.spatial(slide,
                  cmap='viridis',
                  color=mod_sk.location_factors_df.columns,
                  size=1, img_key='hires',
                  alpha_img=0.5, alpha = 0.4, ncols=2,
                  vmin=0, vmax='p99.2'
                 )

Trying to set attribute `.uns` of view, copying.

R = 10

# number of cell type combinations - educated guess assuming that most cell types don't co-locate
n_fact = int(10)

# extract cell abundance from cell2location
X_data = adata_vis.uns['mod']['post_sample_q05']['spot_factors']

import cell2location.models as c2l
# create model class
mod_sk = c2l.CoLocatedGroupsSklearnNMF(n_fact, X_data,
        n_iter = 10000,
        verbose = True,
        var_names=adata_vis.uns['mod']['fact_names'],
        obs_names=adata_vis.obs_names,
        fact_names=['fact_' + str(i) for i in range(n_fact)],
        sample_id=adata_vis.obs['sample'],
        init='random', random_state=0,
        nmf_kwd_args={'tol':0.0001})

# train 5 times to evaluate stability
mod_sk.fit(n=5, n_type='restart')

init_1 - iterations until convergence: 4383
init_2 - iterations until convergence: 4778
init_3 - iterations until convergence: 4217
init_4 - iterations until convergence: 4280
init_5 - iterations until convergence: 3953

## Do some diagnostics
# evaluate stability by comparing trainin restarts

with mpl.rc_context({'figure.figsize': (10, 8)}):
    mod_sk.evaluate_stability('cell_type_factors', align=True)

# evaluate accuracy of the model
mod_sk.compute_expected()
mod_sk.plot_posterior_mu_vs_data()

# extract parameters into DataFrames
mod_sk.sample2df(node_name='nUMI_factors', ct_node_name = 'cell_type_factors')

# export results to scanpy object
adata_vis = mod_sk.annotate_adata(adata_vis) # as columns to .obs
adata_vis = mod_sk.export2adata(adata_vis, slot_name='mod_sklearn') # as a slot in .uns

# print the fraction of cells of each type located to each combination
mod_sk.print_gene_loadings(loadings_attr='cell_type_fractions',
                         gene_fact_name='cell_type_fractions', top_n = 9)

	top-1	top-2	top-3	top-4	top-5	top-6	top-7	top-8	top-9
mean_cell_type_factorsfact_0	IgG+ Plasma: 1.0	Stromal 1: 0.42	Stalk-like ECs: 0.41	Tip-like ECs: 0.4	Stromal 2: 0.38	Mast cells: 0.38	Stromal 3: 0.35	Lymphatic ECs: 0.34	Pericytes: 0.33
mean_cell_type_factorsfact_1	Mature Enterocytes type 1: 0.87	Mature Enterocytes type 2: 0.83	Intermediate: 0.3	IgA+ Plasma: 0.17	Goblet cells: 0.085	Stromal 1: 0.056	gamma delta T cells: 0.051	Stromal 2: 0.047	cDC: 0.037
mean_cell_type_factorsfact_2	CMS3: 0.93	IgA+ Plasma: 0.41	Pro-inflammatory: 0.26	Enteric glial cells: 0.2	Tip-like ECs: 0.11	Stromal 2: 0.092	Lymphatic ECs: 0.072	Stalk-like ECs: 0.072	Mature Enterocytes type 2: 0.072
mean_cell_type_factorsfact_3	CD8+ T cells: 0.64	NK cells: 0.5	Proliferative ECs: 0.49	SPP1+: 0.49	Proliferating: 0.46	gamma delta T cells: 0.23	Unknown: 0.16	T helper 17 cells: 0.16	T follicular helper cells: 0.14
mean_cell_type_factorsfact_4	Stem-like/TA: 0.92	Goblet cells: 0.9	Intermediate: 0.33	Stromal 1: 0.13	Mature Enterocytes type 1: 0.094	Mature Enterocytes type 2: 0.08	IgA+ Plasma: 0.076	CMS3: 0.044	Mast cells: 0.02
mean_cell_type_factorsfact_5	CD19+CD20+ B: 0.94	CD8+ T cells: 0.3	NK cells: 0.21	cDC: 0.17	gamma delta T cells: 0.099	T follicular helper cells: 0.09	T helper 17 cells: 0.083	CMS4: 0.077	Regulatory T cells: 0.051
mean_cell_type_factorsfact_6	CMS2: 0.95	Stromal 2: 0.29	Proliferative ECs: 0.24	Proliferating: 0.24	Pro-inflammatory: 0.21	CMS4: 0.18	Pericytes: 0.12	T follicular helper cells: 0.093	CD4+ T cells: 0.082
mean_cell_type_factorsfact_7	Smooth muscle cells: 1.0	Stromal 3: 0.54	Stalk-like ECs: 0.41	Lymphatic ECs: 0.33	Enteric glial cells: 0.3	Mast cells: 0.25	Tip-like ECs: 0.19	CD4+ T cells: 0.15	Stromal 1: 0.13
mean_cell_type_factorsfact_8	CMS1: 0.94	Unknown: 0.16	Pro-inflammatory: 0.14	CMS4: 0.083	CD8+ T cells: 0.056	NK cells: 0.052	CD19+CD20+ B: 0.035	CMS2: 0.029	Enteric glial cells: 0.015
mean_cell_type_factorsfact_9	Myofibroblasts: 1.0	SPP1+: 0.36	CD4+ T cells: 0.34	Regulatory T cells: 0.32	T follicular helper cells: 0.28	T helper 17 cells: 0.28	CMS4: 0.28	Mast cells: 0.27	Unknown: 0.27

# make nice names
from re import sub
mod_sk.cell_type_fractions.columns = [sub('mean_cell_type_factors', '', i)
                                      for i in mod_sk.cell_type_fractions.columns]

# plot co-occuring cell type combinations
mod_sk.plot_gene_loadings(mod_sk.var_names_read, mod_sk.var_names_read,
                        fact_filt=mod_sk.fact_filt,
                        loadings_attr='cell_type_fractions',
                        gene_fact_name='cell_type_fractions',
                        cmap='RdPu', figsize=[10, 10])

R = 15

# number of cell type combinations - educated guess assuming that most cell types don't co-locate
n_fact = int(15)

# extract cell abundance from cell2location
X_data = adata_vis.uns['mod']['post_sample_q05']['spot_factors']

import cell2location.models as c2l
# create model class
mod_sk = c2l.CoLocatedGroupsSklearnNMF(n_fact, X_data,
        n_iter = 10000,
        verbose = True,
        var_names=adata_vis.uns['mod']['fact_names'],
        obs_names=adata_vis.obs_names,
        fact_names=['fact_' + str(i) for i in range(n_fact)],
        sample_id=adata_vis.obs['sample'],
        init='random', random_state=0,
        nmf_kwd_args={'tol':0.0001})

# train 5 times to evaluate stability
mod_sk.fit(n=5, n_type='restart')

init_1 - iterations until convergence: 3841
init_2 - iterations until convergence: 3708
init_3 - iterations until convergence: 4075
init_4 - iterations until convergence: 3390
init_5 - iterations until convergence: 3354

## Do some diagnostics
# evaluate stability by comparing trainin restarts

with mpl.rc_context({'figure.figsize': (10, 8)}):
    mod_sk.evaluate_stability('cell_type_factors', align=True)

# evaluate accuracy of the model
mod_sk.compute_expected()
mod_sk.plot_posterior_mu_vs_data()

# extract parameters into DataFrames
mod_sk.sample2df(node_name='nUMI_factors', ct_node_name = 'cell_type_factors')

# export results to scanpy object
adata_vis = mod_sk.annotate_adata(adata_vis) # as columns to .obs
adata_vis = mod_sk.export2adata(adata_vis, slot_name='mod_sklearn') # as a slot in .uns

# print the fraction of cells of each type located to each combination
mod_sk.print_gene_loadings(loadings_attr='cell_type_fractions',
                         gene_fact_name='cell_type_fractions', top_n = 9)

	top-1	top-2	top-3	top-4	top-5	top-6	top-7	top-8	top-9
mean_cell_type_factorsfact_0	Mature Enterocytes type 2: 0.81	Mature Enterocytes type 1: 0.22	Intermediate: 0.2	Goblet cells: 0.054	CMS3: 0.029	gamma delta T cells: 0.027	Stem-like/TA: 0.018	Stromal 2: 0.014	IgA+ Plasma: 0.009
mean_cell_type_factorsfact_1	CD8+ T cells: 0.69	NK cells: 0.45	CD19+CD20+ B: 0.31	Proliferating: 0.26	Proliferative ECs: 0.24	gamma delta T cells: 0.16	T follicular helper cells: 0.11	T helper 17 cells: 0.11	Unknown: 0.091
mean_cell_type_factorsfact_2	CD19+CD20+ B: 0.65	CD8+ T cells: 0.23	NK cells: 0.14	cDC: 0.094	T follicular helper cells: 0.061	gamma delta T cells: 0.057	CMS4: 0.053	T helper 17 cells: 0.052	Regulatory T cells: 0.035
mean_cell_type_factorsfact_3	CMS2: 0.95	CMS4: 0.12	Proliferating: 0.12	Pro-inflammatory: 0.11	Proliferative ECs: 0.073	T follicular helper cells: 0.047	CD4+ T cells: 0.037	Intermediate: 0.032	gamma delta T cells: 0.031
mean_cell_type_factorsfact_4	Mature Enterocytes type 1: 0.76	Intermediate: 0.16	Pro-inflammatory: 0.16	Stromal 1: 0.16	cDC: 0.13	gamma delta T cells: 0.11	T helper 17 cells: 0.1	Regulatory T cells: 0.073	Unknown: 0.068
mean_cell_type_factorsfact_5	Goblet cells: 0.83	Intermediate: 0.15	Mature Enterocytes type 2: 0.039	CMS3: 0.036	Stem-like/TA: 0.034	Mature Enterocytes type 1: 0.02	Mast cells: 0.014	gamma delta T cells: 0.013	cDC: 0.0074
mean_cell_type_factorsfact_6	Myofibroblasts: 0.99	CD4+ T cells: 0.21	Regulatory T cells: 0.19	Mast cells: 0.19	T follicular helper cells: 0.18	Unknown: 0.18	T helper 17 cells: 0.16	CMS4: 0.15	gamma delta T cells: 0.15
mean_cell_type_factorsfact_7	Stromal 2: 0.91	Pericytes: 0.81	Proliferative ECs: 0.69	Tip-like ECs: 0.62	Proliferating: 0.54	cDC: 0.32	Pro-inflammatory: 0.32	Lymphatic ECs: 0.31	Regulatory T cells: 0.29
mean_cell_type_factorsfact_8	Smooth muscle cells: 0.99	Stromal 3: 0.52	Stalk-like ECs: 0.32	Enteric glial cells: 0.25	Lymphatic ECs: 0.25	Mast cells: 0.2	CD4+ T cells: 0.1	Tip-like ECs: 0.094	T follicular helper cells: 0.073
mean_cell_type_factorsfact_9	SPP1+: 0.85	cDC: 0.18	CMS4: 0.17	Unknown: 0.076	Stalk-like ECs: 0.073	Regulatory T cells: 0.073	gamma delta T cells: 0.069	Tip-like ECs: 0.063	CD4+ T cells: 0.063
mean_cell_type_factorsfact_10	CMS1: 0.92	Unknown: 0.13	Pro-inflammatory: 0.11	CMS4: 0.067	CD8+ T cells: 0.049	NK cells: 0.044	CD19+CD20+ B: 0.028	CMS2: 0.02	gamma delta T cells: 0.018
mean_cell_type_factorsfact_11	IgA+ Plasma: 0.98	Stromal 1: 0.28	Stalk-like ECs: 0.11	cDC: 0.073	Intermediate: 0.069	Mast cells: 0.059	Tip-like ECs: 0.054	Stromal 2: 0.045	T helper 17 cells: 0.034
mean_cell_type_factorsfact_12	CMS3: 0.91	Enteric glial cells: 0.16	Pro-inflammatory: 0.15	Mast cells: 0.037	Mature Enterocytes type 2: 0.034	Lymphatic ECs: 0.032	Stalk-like ECs: 0.031	Tip-like ECs: 0.031	Goblet cells: 0.019
mean_cell_type_factorsfact_13	Stem-like/TA: 0.94	Intermediate: 0.15	Goblet cells: 0.081	Mature Enterocytes type 2: 0.073	Stromal 1: 0.072	CMS3: 0.026	Stalk-like ECs: 0.016	CD4+ T cells: 0.007	Mast cells: 0.0042
mean_cell_type_factorsfact_14	IgG+ Plasma: 1.0	Stromal 3: 0.36	Stalk-like ECs: 0.28	Mast cells: 0.25	Enteric glial cells: 0.21	Lymphatic ECs: 0.19	Unknown: 0.16	CD4+ T cells: 0.16	Regulatory T cells: 0.14

# make nice names
from re import sub
mod_sk.cell_type_fractions.columns = [sub('mean_cell_type_factors', '', i)
                                      for i in mod_sk.cell_type_fractions.columns]

# plot co-occuring cell type combinations
mod_sk.plot_gene_loadings(mod_sk.var_names_read, mod_sk.var_names_read,
                        fact_filt=mod_sk.fact_filt,
                        loadings_attr='cell_type_fractions',
                        gene_fact_name='cell_type_fractions',
                        cmap='RdPu', figsize=[10, 10])

R = 20

# number of cell type combinations - educated guess assuming that most cell types don't co-locate
n_fact = int(20)

# extract cell abundance from cell2location
X_data = adata_vis.uns['mod']['post_sample_q05']['spot_factors']

import cell2location.models as c2l
# create model class
mod_sk = c2l.CoLocatedGroupsSklearnNMF(n_fact, X_data,
        n_iter = 10000,
        verbose = True,
        var_names=adata_vis.uns['mod']['fact_names'],
        obs_names=adata_vis.obs_names,
        fact_names=['fact_' + str(i) for i in range(n_fact)],
        sample_id=adata_vis.obs['sample'],
        init='random', random_state=0,
        nmf_kwd_args={'tol':0.0001})

# train 5 times to evaluate stability
mod_sk.fit(n=5, n_type='restart')

init_1 - iterations until convergence: 3381
init_2 - iterations until convergence: 3593
init_3 - iterations until convergence: 2618
init_4 - iterations until convergence: 2642
init_5 - iterations until convergence: 4732

## Do some diagnostics
# evaluate stability by comparing trainin restarts

with mpl.rc_context({'figure.figsize': (10, 8)}):
    mod_sk.evaluate_stability('cell_type_factors', align=True)

# evaluate accuracy of the model
mod_sk.compute_expected()
mod_sk.plot_posterior_mu_vs_data()

# extract parameters into DataFrames
mod_sk.sample2df(node_name='nUMI_factors', ct_node_name = 'cell_type_factors')

# export results to scanpy object
adata_vis = mod_sk.annotate_adata(adata_vis) # as columns to .obs
adata_vis = mod_sk.export2adata(adata_vis, slot_name='mod_sklearn') # as a slot in .uns

# print the fraction of cells of each type located to each combination
mod_sk.print_gene_loadings(loadings_attr='cell_type_fractions',
                         gene_fact_name='cell_type_fractions', top_n = 9)

	top-1	top-2	top-3	top-4	top-5	top-6	top-7	top-8	top-9
mean_cell_type_factorsfact_0	Mature Enterocytes type 2: 0.79	Mature Enterocytes type 1: 0.19	Intermediate: 0.18	Goblet cells: 0.049	gamma delta T cells: 0.029	CMS3: 0.028	Stem-like/TA: 0.017	T helper 17 cells: 0.013	cDC: 0.012
mean_cell_type_factorsfact_1	CD8+ T cells: 0.82	NK cells: 0.5	CD19+CD20+ B: 0.3	gamma delta T cells: 0.15	T helper 17 cells: 0.093	T follicular helper cells: 0.087	Proliferative ECs: 0.08	CMS1: 0.071	SPP1+: 0.068
mean_cell_type_factorsfact_2	CD19+CD20+ B: 0.64	cDC: 0.11	CD8+ T cells: 0.044	CMS4: 0.037	T follicular helper cells: 0.034	T helper 17 cells: 0.019	Regulatory T cells: 0.014	NK cells: 0.01	Intermediate: 0.0086
mean_cell_type_factorsfact_3	IgG+ Plasma: 0.93	Mast cells: 0.077	Unknown: 0.076	CD4+ T cells: 0.055	Stromal 1: 0.05	Tip-like ECs: 0.048	Regulatory T cells: 0.035	T follicular helper cells: 0.032	T helper 17 cells: 0.03
mean_cell_type_factorsfact_4	Stem-like/TA: 0.94	Intermediate: 0.14	Goblet cells: 0.079	Mature Enterocytes type 2: 0.075	Stromal 1: 0.066	CMS3: 0.023	Stalk-like ECs: 0.011	Mature Enterocytes type 1: 0.0079	CD4+ T cells: 0.0023
mean_cell_type_factorsfact_5	Mature Enterocytes type 1: 0.78	Intermediate: 0.15	Stromal 1: 0.1	gamma delta T cells: 0.053	Goblet cells: 0.044	CMS4: 0.028	NK cells: 0.026	Mature Enterocytes type 2: 0.023	T helper 17 cells: 0.02
mean_cell_type_factorsfact_6	IgA+ Plasma: 0.97	Stromal 1: 0.24	Stalk-like ECs: 0.058	Intermediate: 0.058	Tip-like ECs: 0.047	cDC: 0.046	Mast cells: 0.037	Mature Enterocytes type 2: 0.022	T helper 17 cells: 0.016
mean_cell_type_factorsfact_7	cDC: 0.61	Regulatory T cells: 0.6	Lymphatic ECs: 0.57	T follicular helper cells: 0.55	T helper 17 cells: 0.5	CD4+ T cells: 0.48	Unknown: 0.47	Stromal 1: 0.43	gamma delta T cells: 0.38
mean_cell_type_factorsfact_8	Smooth muscle cells: 0.88	Lymphatic ECs: 0.091	Stromal 3: 0.077	Enteric glial cells: 0.062	Mast cells: 0.042	Tip-like ECs: 0.033	CD4+ T cells: 0.032	Stromal 1: 0.022	CMS4: 0.019
mean_cell_type_factorsfact_9	CMS3: 0.92	Enteric glial cells: 0.12	Mature Enterocytes type 2: 0.03	Mast cells: 0.028	Lymphatic ECs: 0.021	Goblet cells: 0.021	Tip-like ECs: 0.019	Unknown: 0.0099	IgA+ Plasma: 0.0083
mean_cell_type_factorsfact_10	Pro-inflammatory: 0.82	Proliferative ECs: 0.089	Stalk-like ECs: 0.056	Regulatory T cells: 0.039	Lymphatic ECs: 0.026	T helper 17 cells: 0.014	Unknown: 0.013	NK cells: 0.011	CD4+ T cells: 0.0079
mean_cell_type_factorsfact_11	Proliferating: 0.92	Proliferative ECs: 0.17	T follicular helper cells: 0.046	Pericytes: 0.04	Unknown: 0.038	T helper 17 cells: 0.033	CD4+ T cells: 0.032	Regulatory T cells: 0.029	gamma delta T cells: 0.022
mean_cell_type_factorsfact_12	Pericytes: 0.87	Tip-like ECs: 0.7	Proliferative ECs: 0.56	Stalk-like ECs: 0.19	Lymphatic ECs: 0.059	CD4+ T cells: 0.055	Enteric glial cells: 0.049	NK cells: 0.029	Regulatory T cells: 0.026
mean_cell_type_factorsfact_13	Stromal 2: 0.93	Stromal 1: 0.082	cDC: 0.07	Proliferative ECs: 0.068	CMS4: 0.048	Tip-like ECs: 0.048	Pericytes: 0.035	Mast cells: 0.034	Intermediate: 0.032
mean_cell_type_factorsfact_14	CMS1: 0.93	Unknown: 0.11	CMS4: 0.071	NK cells: 0.034	CD19+CD20+ B: 0.031	CD8+ T cells: 0.025	gamma delta T cells: 0.024	Pro-inflammatory: 0.016	CMS2: 0.014
mean_cell_type_factorsfact_15	SPP1+: 0.83	CMS4: 0.14	cDC: 0.11	Unknown: 0.051	gamma delta T cells: 0.045	NK cells: 0.038	Regulatory T cells: 0.035	T helper 17 cells: 0.031	CD4+ T cells: 0.028
mean_cell_type_factorsfact_16	Stromal 3: 0.86	Stalk-like ECs: 0.65	Enteric glial cells: 0.38	Mast cells: 0.34	Lymphatic ECs: 0.21	CD4+ T cells: 0.081	T follicular helper cells: 0.053	Unknown: 0.046	T helper 17 cells: 0.045
mean_cell_type_factorsfact_17	CMS2: 0.95	CMS4: 0.12	T follicular helper cells: 0.047	Proliferative ECs: 0.041	CD4+ T cells: 0.036	gamma delta T cells: 0.032	Intermediate: 0.031	Regulatory T cells: 0.029	T helper 17 cells: 0.028
mean_cell_type_factorsfact_18	Myofibroblasts: 0.91	CD4+ T cells: 0.11	CMS4: 0.1	Unknown: 0.095	Mast cells: 0.09	gamma delta T cells: 0.087	Regulatory T cells: 0.084	T helper 17 cells: 0.08	T follicular helper cells: 0.078
mean_cell_type_factorsfact_19	Goblet cells: 0.81	Intermediate: 0.14	Mature Enterocytes type 2: 0.038	Stem-like/TA: 0.034	CMS3: 0.031	Mature Enterocytes type 1: 0.015	Mast cells: 0.012	gamma delta T cells: 0.011	T helper 17 cells: 0.0061

# make nice names
from re import sub
mod_sk.cell_type_fractions.columns = [sub('mean_cell_type_factors', '', i)
                                      for i in mod_sk.cell_type_fractions.columns]

# plot co-occuring cell type combinations
mod_sk.plot_gene_loadings(mod_sk.var_names_read, mod_sk.var_names_read,
                        fact_filt=mod_sk.fact_filt,
                        loadings_attr='cell_type_fractions',
                        gene_fact_name='cell_type_fractions',
                        cmap='RdPu', figsize=[10, 10])